from collections.abc import Iterable
from typing import TypeVar

from zkl_aiutils_datasets.basics import Dataset
from zkl_aiutils_datasets.wrapping import wrap_dataset
from .delegating import DelegatingDataset
from .mapped import MappedDataset, SrcSample

AnySample = TypeVar('AnySample')


class PickedDataset(DelegatingDataset[AnySample]):
    def __init__(self, dataset: Dataset[AnySample], indices: Iterable[int]):
        self._dataset = dataset
        self._indices = indices
        super().__init__(MappedDataset(wrap_dataset(indices), lambda index: dataset[index]))

    @property
    def dataset(self) -> Dataset[SrcSample]:
        return self._dataset

    @property
    def indices(self) -> Iterable[int]:
        return self._indices
